/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#ifndef SRC_FLOW_H_
#define SRC_FLOW_H_

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include "FlowDirection.h"
#include "FlowInfo.h"
#include <boost/format.hpp>
#include "Pointer.h"
#include "Packet.h"
#include "IPAddress.h"
#include "ipset/IPAbstractSet.h"
#include "ipset/IPSet.h"
#include "regex/RegexManager.h"
#include "StringCache.h"
#include "L7Payload.h"
#include "protocols/frequency/Frequencies.h"
#include "protocols/frequency/PacketFrequencies.h"
#include "protocols/dns/DNSInfo.h"
#include "protocols/tcp/TCPInfo.h"
#include "protocols/gprs/GPRSInfo.h"
#include "protocols/http/HTTPInfo.h"
#include "protocols/ssl/SSLInfo.h"
#include "protocols/smtp/SMTPInfo.h"
#include "protocols/imap/IMAPInfo.h"
#include "protocols/pop/POPInfo.h"
#include "protocols/sip/SIPInfo.h"
#include "protocols/ssdp/SSDPInfo.h"
#include "protocols/bitcoin/BitcoinInfo.h"
#include "protocols/coap/CoAPInfo.h"
#include "protocols/mqtt/MQTTInfo.h"
#include "protocols/netbios/NetbiosInfo.h"
#include "protocols/dhcp/DHCPInfo.h"
#include "protocols/smb/SMBInfo.h"
#include "protocols/dhcp6/DHCPv6Info.h"
#include "protocols/ssh/SSHInfo.h"
#include "protocols/dcerpc/DCERPCInfo.h"
#include "protocols/quic/QuicInfo.h"
#include "protocols/dtls/DTLSInfo.h"
#include <netinet/in.h>
#include <arpa/inet.h>
#include "AnomalyManager.h"

namespace aiengine {

class TCPInfo;
class FlowForwarder;

class Flow {
public:
    	explicit Flow() { reset(); }
    	explicit Flow(unsigned long hash) { reset(); setId(hash); }
    	virtual ~Flow() {}

	// Common fields of the Flow
	void setId(unsigned long hash) { hash_= hash; }
	unsigned long getId() const { return hash_; }

	// The flow have been marked as reject (RST)
	bool isReject() const { return reject_; }
	void setReject(bool reject) { reject_ = reject; }

	// The flow have been marked to write as evidence
	bool haveEvidence() const { return have_evidence_; }
	void setEvidence(bool value) { have_evidence_ = value; }

	bool isPartialReject() const { return partial_reject_; }
	void setPartialReject(bool reject) { partial_reject_ = reject; }

	// The flow have been marked by a regex for write the packet on a DatabaseAdaptor
	void setWriteMatchedPacket(bool value) { write_matched_packet_ = value; }
	bool isWriteMatchedPacket() const { return write_matched_packet_; }

	// The flow have been analise by the FrequencyEngine
	void setFrequenciesAnalysed(bool value) { frequency_engine_inspected_ = value; }
	bool isFrequenciesAnalysed() const { return frequency_engine_inspected_; }

	bool haveTag() const { return have_tag_; }
	void setTag(uint32_t tag) { have_tag_ = true; tag_ = tag; }
	uint32_t getTag() const { return tag_; }

	void setFlowDirection(FlowDirection dir) { prev_direction_ = direction_; direction_ = dir; }
	FlowDirection getFlowDirection() { return direction_; }
	FlowDirection getPrevFlowDirection() { return prev_direction_; }

	void setPacketAnomaly(const PacketAnomalyType &pa) { pa_ = pa; }
	PacketAnomalyType getPacketAnomaly() const { return pa_; }
	const char *getFlowAnomalyString() const { return PacketAnomalyTypeString[static_cast<std::int8_t>(pa_)].name; }

	// The user label the flow as wanted
	void setLabel(const std::string &label) { label_.reset(); label_ = SharedPointer<std::string>(new std::string(label)); }
	const char *getLabel() const { return  (label_ ? label_->c_str() : ""); }

	// IP functions
	void setFiveTuple(uint32_t src_a, uint16_t src_p, uint16_t proto, uint32_t dst_a, uint16_t dst_p);
       	void setFiveTuple6(struct in6_addr *src_a, uint16_t src_p, uint16_t proto, struct in6_addr *dst_a, uint16_t dst_p);

	// Methods for access to the IP addresses, ports and protocol
	uint32_t getSourceAddress() const { return address_.getSourceAddress(); }
	uint32_t getDestinationAddress() const { return address_.getDestinationAddress(); }
	uint16_t getSourcePort() const { return source_port_; }
	uint16_t getDestinationPort() const { return dest_port_; }
	uint16_t getProtocol() const { return protocol_; }
        const char* getSrcAddrDotNotation() const { return address_.getSrcAddrDotNotation(); }
        const char* getDstAddrDotNotation() const { return address_.getDstAddrDotNotation(); }
	bool isIPv4() const { return (address_.getType() == IPPROTO_IP); }

	struct in6_addr *getSourceAddress6() const { return  address_.getSourceAddress6(); }
	struct in6_addr *getDestinationAddress6() const { return address_.getDestinationAddress6(); }

	const char* status(std::time_t now, int timeout) const;
	const char* getL7ProtocolName() const;

	IPAddress& getAddress() { return address_; }

	// Methods for flow time management
	// TODO: Verify that the current_time_ is allways > than the arrive_time, in some cases
	// conversations could be on different pcap files on reverse time order.
	void setArriveTime(time_t t) { arrive_time_ = t; }
	void setLastPacketTime(time_t t) { current_time_ = t; }

	// For update the flow time on the FlowManager
	struct updateTime {
		explicit updateTime(time_t t): t_(t) {}
		void operator()(const SharedPointer<Flow>& f) {
			f->setLastPacketTime(t_);
		}
		private:
			time_t t_;
	};

	std::time_t getLastPacketTime() const { return current_time_; }
	std::time_t getDuration() const { return (current_time_ - arrive_time_); }
	std::time_t getArriveTime() const { return arrive_time_; }

	uint8_t upstream_ttl;
	uint8_t downstream_ttl;
	uint32_t total_bytes[2];
	uint32_t total_packets[2];
	uint32_t total_packets_l7;
#if defined(BINDING)
	uint32_t total_drop_packets;
	uint32_t total_drop_bytes;
#endif
	// Objects that links with the Flow
	WeakPointer<IPAbstractSet> ipset;
	WeakPointer<Regex> regex;
	WeakPointer<FlowForwarder> forwarder;
	SharedPointer<RegexManager> regex_mng;

	// TCPInfo or GPRSInfo should be attach here
	SharedPointer<FlowInfo> layer4info;

	// Layer 7 objects
	SharedPointer<FlowInfo> layer7info;

	SharedPointer<TCPInfo> getTCPInfo() const { return DynamicPointerCast<TCPInfo>(layer4info); }
	SharedPointer<GPRSInfo> getGPRSInfo() const { return DynamicPointerCast<GPRSInfo>(layer4info); }

        SharedPointer<DNSInfo> getDNSInfo() const { return DynamicPointerCast<DNSInfo>(layer7info); }
        SharedPointer<SSLInfo> getSSLInfo() const { return DynamicPointerCast<SSLInfo>(layer7info); }
        SharedPointer<HTTPInfo> getHTTPInfo() const { return DynamicPointerCast<HTTPInfo>(layer7info); }
        SharedPointer<IMAPInfo> getIMAPInfo() const { return DynamicPointerCast<IMAPInfo>(layer7info); }
        SharedPointer<POPInfo> getPOPInfo() const { return DynamicPointerCast<POPInfo>(layer7info); }
        SharedPointer<SSDPInfo> getSSDPInfo() const { return DynamicPointerCast<SSDPInfo>(layer7info); }
        SharedPointer<SIPInfo> getSIPInfo() const { return DynamicPointerCast<SIPInfo>(layer7info); }
        SharedPointer<SMTPInfo> getSMTPInfo() const { return DynamicPointerCast<SMTPInfo>(layer7info); }
	SharedPointer<BitcoinInfo> getBitcoinInfo() const { return DynamicPointerCast<BitcoinInfo>(layer7info); }
	SharedPointer<CoAPInfo> getCoAPInfo() const { return DynamicPointerCast<CoAPInfo>(layer7info); }
	SharedPointer<MQTTInfo> getMQTTInfo() const { return DynamicPointerCast<MQTTInfo>(layer7info); }
	SharedPointer<NetbiosInfo> getNetbiosInfo() const { return DynamicPointerCast<NetbiosInfo>(layer7info); }
	SharedPointer<DHCPInfo> getDHCPInfo() const { return DynamicPointerCast<DHCPInfo>(layer7info); }
	SharedPointer<SMBInfo> getSMBInfo() const { return DynamicPointerCast<SMBInfo>(layer7info); }
	SharedPointer<DHCPv6Info> getDHCPv6Info() const { return DynamicPointerCast<DHCPv6Info>(layer7info); }
	SharedPointer<SSHInfo> getSSHInfo() const { return DynamicPointerCast<SSHInfo>(layer7info); }
	SharedPointer<DCERPCInfo> getDCERPCInfo() const { return DynamicPointerCast<DCERPCInfo>(layer7info); }
	SharedPointer<QuicInfo> getQuicInfo() const { return DynamicPointerCast<QuicInfo>(layer7info); }
	SharedPointer<DTLSInfo> getDTLSInfo() const { return DynamicPointerCast<DTLSInfo>(layer7info); }

	// Special objects for frequency analisys
	SharedPointer<Frequencies> frequencies;
	SharedPointer<PacketFrequencies> packet_frequencies;

	Packet *packet;

	// specific values for a specific Engine

	void reset();

	friend std::ostream& operator<< (std::ostream &out, const Flow &flow);

	void show(std::ostream &out) const;
	void showL7(std::ostream &out) const;
	void showL4(std::ostream &out) const;
	void show(Json &out) const;

#if defined(BINDING)

	// The user over the Rest API wants to copy and see the latest L7 payload
	SharedPointer<L7Payload> l7_payload;

	// By doing this operation the flow is reevaluated with other
	// FlowForwarders so could be link to the protocol, this is usefull
	// for moving flows from tcpgeneric/udpgeneric to their protocol.
	void detach();

	// The flow have been marked as accept or drop (for external Firewall integration)
	bool isAccept() const { return is_accept_; }
	void setAccept(bool accept) { is_accept_ = accept; }

	bool isAlerted() const { return is_alerted_; }
	void setAlerted(bool value) { is_alerted_ = value; }

	const char* getDirection() const;
	uint32_t getTotalUpstreamBytes() const { return total_bytes[static_cast<int>(FlowDirection::FORWARD)]; }
	uint32_t getTotalDownstreamBytes() const { return total_bytes[static_cast<int>(FlowDirection::BACKWARD)]; }
	uint32_t getTotalUpstreamPackets() const { return total_packets[static_cast<int>(FlowDirection::FORWARD)]; }
	uint32_t getTotalDownstreamPackets() const { return total_packets[static_cast<int>(FlowDirection::BACKWARD)]; }
        uint32_t getTotalBytes() const { return getTotalUpstreamBytes() + getTotalDownstreamBytes(); }
        uint32_t getTotalPacketsLayer7() const { return total_packets_l7; }
        uint32_t getTotalPackets() const { return getTotalUpstreamPackets() + getTotalDownstreamPackets(); }
	uint8_t getUpstreamTTL() const { return upstream_ttl; }
	uint8_t getDownstreamTTL() const { return downstream_ttl; }

        Frequencies* getFrequencies() const { return frequencies.get(); }
        PacketFrequencies* getPacketFrequencies() const { return packet_frequencies.get(); }

        Regex* getRegex() const { return regex.lock().get(); }
	RegexManager *getRegexManager() const { return regex_mng.get(); }

	TCPInfo* getTCPInfoObject() const { return getTCPInfo().get(); }

	IPAbstractSet* getIPSetInfo() const { return ipset.lock().get(); }
        HTTPInfo* getHTTPInfoObject() const { return getHTTPInfo().get(); }
        SIPInfo* getSIPInfoObject() const { return getSIPInfo().get(); }
        DNSInfo* getDNSInfoObject() const { return getDNSInfo().get(); }
        SSLInfo* getSSLInfoObject() const { return getSSLInfo().get(); }
        SMTPInfo* getSMTPInfoObject() const { return getSMTPInfo().get(); }
        POPInfo* getPOPInfoObject() const { return getPOPInfo().get(); }
        IMAPInfo* getIMAPInfoObject() const { return getIMAPInfo().get(); }
        SSDPInfo* getSSDPInfoObject() const { return getSSDPInfo().get(); }
	BitcoinInfo* getBitcoinInfoObject() const { return getBitcoinInfo().get(); }
        CoAPInfo* getCoAPInfoObject() const { return getCoAPInfo().get(); }
        MQTTInfo* getMQTTInfoObject() const { return getMQTTInfo().get(); }
        NetbiosInfo* getNetbiosInfoObject() const { return getNetbiosInfo().get(); }
        DHCPInfo* getDHCPInfoObject() const { return getDHCPInfo().get(); }
        SMBInfo* getSMBInfoObject() const { return getSMBInfo().get(); }
        DHCPv6Info* getDHCPv6InfoObject() const { return getDHCPv6Info().get(); }
        SSHInfo* getSSHInfoObject() const { return getSSHInfo().get(); }
        DCERPCInfo* getDCERPCInfoObject() const { return getDCERPCInfo().get(); }
        QuicInfo* getQuicInfoObject() const { return getQuicInfo().get(); }
        DTLSInfo* getDTLSInfoObject() const { return getDTLSInfo().get(); }
#endif

#if defined(PYTHON_BINDING)
	boost::python::list getPayload() const;
	void setRegexManager(const SharedPointer<RegexManager> &rm);
#elif defined(RUBY_BINDING)
	VALUE getPayload() const;
#elif defined(LUA_BINDING)
	RawPacket& getPacket() const;
	const char *__str__();
#elif defined(JAVA_BINDING)
	IPAbstractSet& getIPSet() const;
#endif

private:
	unsigned long hash_;
	IPAddress address_;
	uint16_t source_port_;
	uint16_t dest_port_;
	uint16_t protocol_;
	uint32_t tag_;
	bool have_tag_:1;
	bool reject_:1; // The flow can be reject from the ruby/python side
	bool partial_reject_:1; // For UDP flows
	bool have_evidence_:1;
#if defined(BINDING)
	bool is_accept_:1;
	bool is_alerted_:1;
#endif
	mutable bool write_matched_packet_:1;
	bool frequency_engine_inspected_:1;
	FlowDirection direction_;
	FlowDirection prev_direction_;
	PacketAnomalyType pa_;
	time_t arrive_time_;
	time_t current_time_;
	SharedPointer<std::string> label_;
};

} // namespace aiengine

#endif  // SRC_FLOW_H_
